/***
 * 堆 优先队列
 * 41. 数据流中的中位数
 * 40. 最小的k个数
 * 
 * **/

#include <iostream>
#include <queue>
#include <vector>
using namespace std;

template <class T>
void print(const T &t)
{
    typename T::const_iterator it = t.cbegin();
    for (; it != t.cend(); ++it) {
        cout << *it << " ";
    }
    cout << endl;
}

// 41.
class MedianFinder
{
public:
    void addNum(int num)
    {
        if (_pqMin.empty() || num < _pqMin.top()) {
            _pqMin.push(num);
            if (_pqMin.size() - _pqMax.size() > 1) {
                _pqMax.push(_pqMin.top());
                _pqMin.pop();
            }
        } else {
            _pqMax.push(num);
            if (_pqMax.size() - _pqMin.size() > 1) {
                _pqMin.push(_pqMax.top());
                _pqMax.pop();
            }
        }
    }

    double findMedian()
    {
        if (size() % 2 == 0) {
            return (double)(_pqMin.top() + _pqMax.top()) / 2;
        }
        return _pqMin.size() > _pqMax.size() ? _pqMin.top() : _pqMax.top();
    }

    size_t size() const
    {
        return _pqMax.size() + _pqMin.size();
    }

private:
    priority_queue<int> _pqMin;                                 // 大顶堆
    priority_queue<int, vector<int>, std::greater<int>> _pqMax; // 小顶堆
};

class Solution
{
public:
    // 40. topK 基于size==k的堆
#if 0
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        // 维护size为k的大顶堆
        if (k == 0 || arr.empty()) return {};
        vector<int> res(arr.begin(), arr.begin() + k);
        buildHeap(res, k);  // O(K)
        for (int i = k; i < arr.size(); ++i){
            if (arr[i] < res.front()){
                res[0] = arr[i];
                heapify(res, 0, k);
            }
        }
        return res;
    }
private:
    void heapify(vector<int> &nums, int dad, int length){
        int son = 2 * dad + 1;
        while (son < length){
            if (son + 1 < length && nums[son] < nums[son + 1]){
                ++son;
            }
            if (nums[son] < nums[dad]){
                break;
            }
            swap(nums[son], nums[dad]);
            dad = son;
            son = 2 * dad + 1;
        }
    }
    void buildHeap(vector<int> &nums, int length){
        for (int i = length / 2 - 1; i >= 0; --i){
            heapify(nums, i, length);
        }
    }
#endif
    // 40.topK 基于partition，pivot==k时，停止递归
    vector<int> getLeastNumbers(vector<int> &arr, int k)
    {
        // 利用partition思想，当partition==k时，意味着前半部分就是topK
        if (arr.empty() || k == 0) {
            return {};
        }
        if (arr.size() == k) {
            return vector<int>(arr);
        }
        topK(arr, k, 0, arr.size() - 1);
        vector<int> res(arr.begin(), arr.begin() + k);
        return res;
    }

private:
    int partition(vector<int> &nums, int left, int right)
    {
        int mark = nums[left];
        int pivot = left;
        for (int i = left; i <= right; ++i) {
            if (nums[i] < mark) {
                ++pivot;
                swap(nums[i], nums[pivot]);
            }
        }
        swap(nums[left], nums[pivot]);
        return pivot;
    }

    void topK(vector<int> &nums, int k, int left, int right)
    {
        int pos = partition(nums, left, right);
        if (pos == k) {
            return;
        } else if (pos < k) {
            topK(nums, k, pos + 1, right);
        } else {
            topK(nums, k, left, pos);
        }
    }
};

Solution solve = Solution();

void test40()
{
    vector<int> nums = {0, 0, 2, 3, 2, 1, 1, 2, 0, 4};
    print(solve.getLeastNumbers(nums, 10));
}
int main()
{
    test40();
    return 0;
}